'''
https://leetcode.cn/problems/rearrange-string-k-distance-apart
'''
import heapq
from collections import Counter


class Solution:
    def rearrangeString(self, s: str, k: int) -> str:
        counter = Counter(s)
        pq = []     # (剩余数量(负数，选择最多的),下一次可以出现的最早位置,  字符)
        for c, cnt in counter.items():
            heapq.heappush(pq, (-cnt, 0, c))

        res = []
        while pq:
            temp = []
            i = len(res)
            while pq and i < pq[0][1]:
                temp.append(heapq.heappop(pq))
            if not pq: return ''
            cnt, _ , c = heapq.heappop(pq)
            cnt = -cnt
            res.append(c)
            if cnt - 1:
                heapq.heappush(pq, (-(cnt-1), i+k, c))
            for e in temp:
                heapq.heappush(pq,e)
        return ''.join(res)

s = 'aaabc'
k = 2
print(Solution().rearrangeString(s, k))